import time

from selenium.webdriver.common.by import By
from selenium import webdriver
from selenium.webdriver.common.action_chains import ActionChains
from selenium.webdriver.common.keys import Keys
from selenium.webdriver.firefox.options import Options
from selenium.webdriver.chrome.service import Service

from selenium.webdriver.chrome.options import Options as ChromOptions


class ChatgptSelenium():
    def __init__(self,chatgpt_url,driver_path,local_port=5001,refesh=True):
        '''
        chatgpt 调用网页读写内容
        :param chatgpt_url: chatgpt网站地址
        :param local_port: chrome 本地调试端口
        :param driver_path: chrome webdriver路径
        '''
        self.chatgpt_url = chatgpt_url
        self.browser = self.get_browser(chromedirver_path=driver_path, local_port=local_port)
        self.browser.get(chatgpt_url)
        time.sleep(3)
        if 'qianxian' in self.chatgpt_url:
            self.qianxianwanjia_click()
        time.sleep(2)

        self.textarea_input_element = self.find_textarea_element()
        self.submit_button_element = self.find_submit_element(init=True)
        # self.init_start_index()
        self.refesh = refesh

    def get_browser(self, chromedirver_path, local_port):
        '''
        初始话webdriver browser
        :param chromedirver_path:
        :param local_port:
        :return:
        '''

        if "chatglm" not in self.chatgpt_url:  # and "yiyan" not in self.chatgpt_url:
            options = ChromOptions()

            options.add_argument(r'--user-data-dir=E:\chrome_data\xulaoshi')
            options.add_experimental_option("debuggerAddress", "127.0.0.1:" + str(local_port))
            options.add_argument(
                "user-agent=Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/2blabla")
            browser = webdriver.Chrome(chrome_options=options, executable_path=chromedirver_path)
        else:
            options = Options()
            browser = webdriver.Firefox(options=options, firefox_profile=r"F:\FireFoxNew")
            browser.implicitly_wait(3)
        return browser

    def find_textarea_element(self):
        '''
        找到chatgpt网页的textarea
        :return:
        '''
        same_xpath_url = ["chat.zhile.io", "chat.huiyan-ai.com", "bettergpt.chat", "poe.com"]
        textarea_input_element = None
        for same_url in same_xpath_url:
            if same_url in self.chatgpt_url:
                textarea_input_element = self.browser.find_element(By.XPATH, "//textarea")
                break

        if not textarea_input_element:
            if 'chat.forefront.ai' in self.chatgpt_url:
                textarea_input_element = self.browser.find_element(By.XPATH,
                                                                   '//div[@class="text-th-primary-dark"]/div/span/span')

            elif "promptknit.com" in self.chatgpt_url:
                textarea_input_element = self.browser.find_elements(By.XPATH, "//textarea")[1]
            elif 'yiyan.baidu.com' in self.chatgpt_url:
                textarea_input_element = self.browser.find_element(By.XPATH, '//div[@class="JRWXI4iU"]/textarea')
            elif 'chatglm' in self.chatgpt_url:
                textarea_input_element = self.browser.find_element(By.XPATH, '//div[@class="input-box-inner"]/textarea')
            elif 'baichuan-ai.com' in self.chatgpt_url:
                textarea_input_element = self.browser.find_element(By.XPATH,
                                                                   '//textarea[contains(@class,"chat-input")]')
            elif "achuanai" in self.chatgpt_url or "chat18" in self.chatgpt_url:
                textarea_input_element = self.browser.find_element(By.XPATH,'//textarea[@class="n-input__textarea-el"]')
            elif "you.com" in self.chatgpt_url:
                textarea_input_element = self.browser.find_element(By.XPATH,
                                                                   '//textarea[@id="search-input-textarea"]')
            else:
                raise ValueError("not found")

        return textarea_input_element

    def find_reset_element(self):
        '''
        找到chatgpt网页的重置按钮
        :return:
        '''
        if "poe.com" in self.chatgpt_url:
            reset_button = self.browser.find_element(By.XPATH,
                                                     '//button[contains(@class,"ChatMessageInputFooter_chatBreakButton__hqJ3v")]')
        else:
            raise ValueError("not found")
        return reset_button

    def find_submit_element(self, init=False):
        if "chat.zhile.io" in self.chatgpt_url:
            submit_button = self.browser.find_element(By.XPATH, '//main//form/div/div/button')
        elif "chat.huiyan-ai.com" in self.chatgpt_url:
            submit_button = self.browser.find_element(By.XPATH, '//footer/div/div/button')
        elif "bettergpt.chat" in self.chatgpt_url:
            submit_button = self.browser.find_element(By.XPATH, '//div[@class="flex"]/div/button')
        elif "poe.com" in self.chatgpt_url:
            submit_button = self.browser.find_element(By.XPATH,
                                                      '//button[contains(@class,"ChatMessageSendButton_sendButton__4ZyI4")]')
        elif "chat.forefront.ai" in self.chatgpt_url:
            try:
                submit_button = self.browser.find_element(By.XPATH, '//div[contains(@class,"cursor-pointer self-end")]')
            except Exception as e:
                print("find not submit button")
                submit_button = None
        elif "promptknit.com" in self.chatgpt_url:
            submit_button = None
            submit_buttons = self.browser.find_elements(By.XPATH, '//button')
            for button in submit_buttons:
                if button.text == "Run":
                    submit_button = button
                    break
        elif 'yiyan.baidu.com' in self.chatgpt_url:
            submit_button = self.browser.find_element(By.XPATH, '//span[@class="VAtmtpqL"]')
        elif 'chatglm' in self.chatgpt_url:
            if init:
                submit_button = self.browser.find_element(By.XPATH, '//div[@class="enter"]/div/img')
            else:
                submit_button = self.submit_button_element
        elif 'baichuan-ai.com' in self.chatgpt_url:
            submit_button = self.browser.find_element(By.XPATH, '//div[contains(@class,"send-btn")]')
        elif "achuanai" in self.chatgpt_url:
            submit_button = self.browser.find_element(By.XPATH, '//footer//button[contains(@class,"n-button n-button--default-type")]')
        elif "chat18" in self.chatgpt_url:
            submit_button = self.browser.find_element(By.XPATH, '//footer//button[contains(@class,"n-button--primary-type")]')
        elif "you.com" in self.chatgpt_url:
            submit_button = self.browser.find_element(By.XPATH,'//div[@class="sc-bd4e4c41-1 dRkSSg"]')
        else:
            raise ValueError("not found")
        return submit_button

    def get_chatgpt_result(self, return_all=False):
        '''

        :param text:
        :return:
        '''
        time.sleep(5)
        if "chat.zhile.io" in self.chatgpt_url:
            result_div_elements = self.browser.find_elements(By.XPATH,
                                                             "//main/div/div/div/div/div/div[1]/div/div[1]/div[2]/div[2]")  # '//main/div/div/div/div/div[1]/div[1]/div')
        elif "chat.huiyan-ai.com" in self.chatgpt_url:
            result_div_elements = self.browser.find_elements(By.XPATH, '//div[@class="leading-relaxed break-words"]')
        elif "bettergpt.chat" in self.chatgpt_url:
            result_div_elements = self.browser.find_elements(By.XPATH,
                                                             '//*[@id="root"]/div/div[3]/main/div/div/div/div/div[1]/div/div/div[2]')
        elif "poe.com" in self.chatgpt_url:
            result_div_elements = self.browser.find_elements(By.XPATH,
                                                             '//div[@class="Markdown_markdownContainer__Tz3HQ"]')
        elif "chat.forefront.ai" in self.chatgpt_url:
            result_div_elements = self.browser.find_elements(By.XPATH,
                                                             '//div[contains(@class,"text-base")]')
        elif "promptknit.com" in self.chatgpt_url:
            result_div_elements = self.browser.find_elements(By.XPATH, '//div[contains(@class,"break-words")]')
        elif 'yiyan.baidu.com' in self.chatgpt_url:
            result_div_elements = self.browser.find_elements(By.XPATH,
                                                             '//div[@class="ErXhAgf5"]//div[contains(@class,"custom-html")]')
            result_div_elements = result_div_elements[::-1]
        elif 'chatglm' in self.chatgpt_url:
            result_div_elements = self.browser.find_elements(By.XPATH,
                                                             '//div[@class="item"]/div[contains(@class,"answer")]')
        elif 'baichuan-ai' in self.chatgpt_url:
            result_div_elements = self.browser.find_elements(By.XPATH, '//div[contains(@class,"chat-row-text")]')
        elif 'achuanai' in self.chatgpt_url:
            result_div_elements = self.browser.find_elements(By.XPATH, '//div[@id="message-wrapper"]//div[contains(@class,"leading-relaxed")]')
        elif "you.com" in self.chatgpt_url:
            roots = self.browser.find_elements(By.XPATH,"//div[@id='chatHistory']/div/div")
            result_div_elements = []
            for i in roots:
                child_divs = i.find_elements(By.XPATH,"./div")
                if len(child_divs) >= 2:
                    result_div_elements.append(child_divs[1])
        elif "chat18" in self.chatgpt_url:
            result_div_elements = self.browser.find_elements(By.XPATH,
                                                             '//div[@id="image-wrapper"]//div[contains(@class,"items-end")]')

        else:
            raise ValueError("not found")
        result_div_text = [element.text for element in result_div_elements]
        if result_div_text:
            if return_all:
                return result_div_text
            else:
                return result_div_text[-1]
        else:
            return ""



    def qianxianwanjia_click(self):
        pass
        # time.sleep(5)
        # label_buttons = self.browser.find_elements(By.XPATH,'//main/div/div/div/div/div')
        # label_buttons[1].click()
        # time.sleep(2)

    def send_text_with_huanhang(self, text):


        self.textarea_input_element = self.find_textarea_element()

        self.textarea_input_element.click()
        for sub_text in text.split("\n"):
            self.textarea_input_element.send_keys(sub_text)
            ActionChains(self.browser).key_down(Keys.SHIFT).key_down(Keys.ENTER).key_up(Keys.SHIFT).key_up(
                Keys.ENTER).perform()
            time.sleep(1)

    def cancel_generating(self):
        if "promptknit.com" in self.chatgpt_url:
            buttons = self.browser.find_elements(By.XPATH, '//button')
            for button in buttons:
                if button.text == "Stop Running":
                    button.click()
                    break

    def send_message(self, text, time_sleep=1, prompt_text="", ignore_texts=None):
        '''
        向chatgpt网页发送内容
        :param text:
        :param prompt_text:相关任务提示短语
        :param ignore_texts:需要忽略得文字
        :return:
        '''
        time.sleep(1)
        self.find_textarea_element().clear()
        self.send_text_with_huanhang(text + prompt_text)
        print("sent over")
        time.sleep(1)

        self.find_submit_element().click()
        time.sleep(2)
        checkCount = 0
        while self.check_is_generating() and checkCount < 40:
            time.sleep(3)
            checkCount += 1

        if checkCount >= 40:
            self.cancel_generating()
            self.find_textarea_element().clear()
            return ""

        time.sleep(time_sleep)
        text = self.get_chatgpt_result()
        if ignore_texts:
            for t in ignore_texts:
                text = text.replace(t, "")

        self.find_textarea_element().clear()
        return text

    def get_prompt_result(self, text, time_sleep=1, prompt_text="", ignore_texts=None, reset_click=False):
        success = False
        result = ""

        if reset_click:
            self.find_reset_element().click()
            time.sleep(1)
        self.find_textarea_element().clear()
        while success == False:
            self.find_textarea_element().clear()
            result = self.send_message(text, time_sleep, prompt_text, ignore_texts)
            if result != "":
                if "experiencing heavy traffic" in result:
                    success = False
                elif "Something went wrong" in result:
                    success = False
                elif "No response" in result:
                    success = False
                elif "服务异常" in result:
                    success = False
                elif "AI思考中" in result:
                    success = False
                elif "请稍后再试" in result:
                    success = False
                elif "Internal" in result:
                    success = False
                elif "服务端错误" in result:
                    success = False
                else:
                    success = True
            if not success:
                time.sleep(15)

        return result

    def check_is_generating(self):

        if "chat.zhile.io" in self.chatgpt_url:
            generate_button_text = self.browser.find_element(By.XPATH, "//form").text
            if "Regenerate response" == generate_button_text:
                return 0
            else:
                return 1


        elif "chat.huiyan-ai.com" in self.chatgpt_url:
            generate_button = self.browser.find_elements(By.XPATH, "//span[@class='n-button__content']")
            try:
                button_texts = [element.text for element in generate_button if element.text != ""]
                for text in button_texts:
                    if "Responding" in text:
                        return 1
            except Exception as e:
                return 0
        elif "bettergpt.chat" in self.chatgpt_url:
            status = self.submit_button_element.get_attribute("class")
            if "cursor-not-allowed" in status:
                return 1
            else:
                return 0
        elif "poe.com" in self.chatgpt_url:
            status = self.find_submit_element().is_enabled()
            try:
                last_message = self.get_chatgpt_result()
            except Exception as e:
                last_message = "..."
            if status:
                if last_message == '...':
                    return 1
                else:
                    return 0
            else:
                return 1
        elif "chat.forefront.ai" in self.chatgpt_url:
            generator_div = self.browser.find_element(By.XPATH, '//div[@class="flex gap-3"]')
            if "generat" in generator_div.text:
                return 1
            else:
                return 0

        elif "promptknit.com" in self.chatgpt_url:
            buttons = self.browser.find_elements(By.XPATH, '//button')
            for button in buttons:
                if button.text == "Stop Running":
                    return 1
        elif 'yiyan.baidu.com' in self.chatgpt_url:
            try:
                generate_button = self.browser.find_element(By.XPATH, '//span[@class="mEKFkIX7"]')
                if generate_button and generate_button.text == "停止生成":
                    return 1
            except Exception as e:
                return 0
        elif 'chatglm' in self.chatgpt_url:
            current_text = self.browser.find_element(By.XPATH, '//div[@id="session-container"]').text
            if "停止生成" in current_text:
                return 1
        elif 'baichuan-ai' in self.chatgpt_url:
            try:
                re_generate_text = self.browser.find_element(By.XPATH, '//div[@class="stop-retry-btn"]')
                if re_generate_text and re_generate_text.text == "重新生成":
                    return 0
                else:
                    return 1
            except Exception as e:
                return 1
        elif "achuanai" in self.chatgpt_url or "chat18" in self.chatgpt_url:
            try:
                time.sleep(1)
                re_generate_text = self.browser.find_elements(By.XPATH, '//span[@class="n-button__content"]')
                for re_button in  re_generate_text:
                    if re_button and ("终止生成" in re_button.text or "Stop" in re_button.text):
                        return 1

                return 0
            except Exception as e:
                return 0
        elif "you.com" in self.chatgpt_url:
            try:
                time.sleep(2)
                re_generate_text = self.browser.find_elements(By.XPATH, '//div[contains(@class,"chat-input")]')
                for re_button in  re_generate_text:
                    if re_button and "generating" in re_button.text:
                        return 1

                return 0
            except Exception as e:
                return 0

        return 0